#!/usr/bin/python
'''
    Filename - ftpc.py
    Author: Gerrit Kitts <gkitts@gmail.com>    
    Usage: ftpc.py (remoteIP) (remotePort) (trollPort) (file-to-transfer)
    
    This package provides a very simple client that will transfer binary files to a host
    that is running it's sister process, ftps.py.  Communication is via UDP, and passes
    through a 'troll' process on the localhost to simulate nework package dropping.  
    Each UDP packet consists of:
        remote IP <4 bytes>: The destination IP for the datagram
        remote Port <2 bytes>: The destination port for the datagram
        datagram type <1 byte>: 1 if the datagram contains the file size, 2 if the datagram
            contains the file name, 3 for data segments.
'''

import hashlib
import os, os.path
import select
import socket
import struct
import time

from sys import argv, stderr

FLAG_SIZE = 1
FLAG_NAME = 2
FLAG_DATA = 3

SEQ_ZERO = 0
SEQ_ONE = 1

BUFF_SIZE = 1025
DATA_SIZE = 1000
HEADER_STRUCT = '! 4B H B B' # common header elements for packer functions
SOCKET_TIMEOUT = 0.05 # socket timeout expressed in seconds; is 50ms

'''
    Exception sub-class to handle invalid argument specifications
'''
class ArgumentException(Exception):
    def __init__(self):
       self.message = "Usage: ftps.py (remoteIP) (remotePort) (file-to-transfer)"

''' 
    openFile - Attempts to open the file provided as params['file'], and reads 
      the size and its basename.  
      Return: Returns a tuple containing the file handle, the size of the file 
        and the basename of the file
      Exceptions: If the file does not exist, raises an IOError.      
'''
def openFile(params):
    filename = params['file']
    if not os.path.exists(filename):
        raise IOError("File %s was not found!" % (filename))
    stats = os.stat(filename)
    size = stats.st_size    
    f = open(filename, 'rb')
    basename = os.path.basename(filename)
    return (f, size, basename)


'''
    Decodes the return packet, reading the headers
'''
def decodeHeaders(data):
    unpacker = struct.Struct(HEADER_STRUCT)
    header = unpacker.unpack(data[0:8])
    destIP = '.'.join([str(x) for x in header[0:4]])
    destPort = header[4]
    flag = header[5]
    seq = header[6]
    return (flag, seq)

'''
    Handles the flow control for the send
'''
def sendDatagram(socket, datagram, trollAddr, dgramSeq):
    s = socket
    bytesSent = s.sendto(datagram, trollAddr)
    
    # Wait for the response
    done = False
    retransmit = False
    while not done:
        read, write, err = select.select([s], [], [], SOCKET_TIMEOUT)
        if (len(read) > 0):
            # Received an ACK, hopefully
            data = read[0].recv(BUFF_SIZE)
            (flag, seq) = decodeHeaders(data)
            if seq is dgramSeq:
                # Good read, return
                print "got ACK, continuing"
                retransmit = False
                done = True
            else:
                print "Invalid ACK of %i, resending datagram." % seq
                retransmit = True
        else:
            print ".",
            #print " Timeout, resending datagram."
            retransmit = True
                
        if retransmit:
            # Resend the packet, either due to timeout or ACK failure
            bytesSent = s.sendto(datagram, trollAddr)

'''
    Sends the first datagram indicating file size
'''
def sendFileSizeDatagram(remIP, remPort, size, trollAddr, socket):
    packer = struct.Struct(HEADER_STRUCT + ' I')
    datagram = packer.pack(*(remIP + [remPort, FLAG_SIZE, SEQ_ZERO, size]))
    print "Sending size %i to %s:%i ." % (size, remIP, remPort),
    sendDatagram(socket, datagram, trollAddr, SEQ_ZERO)
            
'''
    Formats the datagram for a name segment
'''
def sendFileNameDatagram(remIP, remPort, name, trollAddr, socket):
    packer = struct.Struct(HEADER_STRUCT + ' 20s')
    datagram = packer.pack(*(remIP + [remPort, FLAG_NAME, SEQ_ONE, name]))
    print "Sending name %s to %s:%i ." % (name, remIP, remPort),
    sendDatagram(socket, datagram, trollAddr, SEQ_ONE)

'''
    Formats the datagram for a data packet
'''
def sendDataDatagram(remIP, remPort, bytes, seq, trollAddr, socket):
    packer = struct.Struct(HEADER_STRUCT + ' 1000s')
    datagram = packer.pack(*(remIP + [remPort, FLAG_DATA, seq, bytes]))    
    print "Sending data with seq %i to %s:%i ." % (seq, remIP, remPort),
    sendDatagram(socket, datagram, trollAddr, seq)

'''
    Attempt to transfer the file to the server specified.  Reads
        params['address'] as the destination, and transfers the file
        in 1000 byte chunks.  
'''
def sendFile(params):
    fileProps = openFile(params)
    print "Sending file %s with size %i" % (fileProps[2], fileProps[1])
    
    # Create our UDP socket, and begin listening / transmitting
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    
    remIP = params['address'][0]
    remPort = params['address'][1]
    print "Remote IP is %s:%s" % (remIP, remPort)
    remIP = [int(x) for x in remIP.split('.')]
    
    trollAddr = ('127.0.0.1', params['trollPort'])
    print "Troll IP is %s:%s" % (trollAddr[0], trollAddr[1])
    
    sendFileSizeDatagram(remIP, remPort, fileProps[1], trollAddr, s)
    sendFileNameDatagram(remIP, remPort, fileProps[2], trollAddr, s)
    
    # Now send the contents of the file in 1000 byte chunks.  For our purposes,
    #  a character is equivalent to a byte so we can read the file directly.    
    m = hashlib.md5()
    f = fileProps[0]    
    bytes = f.read(1000)
    m.update(bytes)
    seq = SEQ_ZERO
    while len(bytes) > 0:
        dataDatagram = sendDataDatagram(remIP, remPort, bytes, seq, trollAddr, s)
        bytes = f.read(1000)
        m.update(bytes[0:len(bytes)])
        
        # Update the sequence number
        if seq is SEQ_ONE:
            seq = SEQ_ZERO
        else:
            seq = SEQ_ONE
    
    # Close the socket to clean up after ourselves
    s.close()
    
    # Close our file handle
    f.close()

    # Finally, determine the MD5 hash for easy comparison
    print "MD5 checksum is: %s" % (m.hexdigest())

'''
    Base wrapper for class functionality
'''
if __name__ == '__main__':
    params = {
        'address' : ('127.0.0.1', 7779),
        'trollPort' : 7775,
        'file' : 'foo.jpg'
    }
    
    if (len(argv) is not 5):
        # Exception case
        raise ArgumentException 
    else:
        params['address'] = (argv[1], int(argv[2]))
        params['trollPort'] = (int(argv[3]))
        params['file'] = argv[4]

    sendFile(params)
